/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.extensions.sql.impl.utils;



import com.bff.gaia.unified.sdk.schemas.Schema;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.BiMap;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableBiMap;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableMap;

import org.apache.calcite.avatica.util.ByteString;

import org.apache.calcite.rel.type.RelDataType;

import org.apache.calcite.rel.type.RelDataTypeFactory;

import org.apache.calcite.rel.type.RelDataTypeField;

import org.apache.calcite.sql.type.SqlTypeName;

import com.bff.gaia.unified.sdk.schemas.LogicalTypes;

import org.joda.time.Instant;

import org.joda.time.base.AbstractInstant;



import java.lang.reflect.Type;

import java.util.Date;

import java.util.Map;

import java.util.stream.IntStream;



/** Utility methods for Calcite related operations. */

public class CalciteUtils {

  private static final long UNLIMITED_ARRAY_SIZE = -1L;



  // SQL has schema types that do not directly correspond to Unified Schema types. We define

  // LogicalTypes to represent each of these types.



  /** A LogicalType corresponding to DATE. */

  public static class DateType extends LogicalTypes.PassThroughLogicalType<Instant> {

    public static final String IDENTIFIER = "SqlDateType";



    public DateType() {

      super(IDENTIFIER, "", Schema.FieldType.DATETIME);

    }

  }



  /** A LogicalType corresponding to TIME. */

  public static class TimeType extends LogicalTypes.PassThroughLogicalType<Instant> {

    public static final String IDENTIFIER = "SqlTimeType";



    public TimeType() {

      super(IDENTIFIER, "", Schema.FieldType.DATETIME);

    }

  }



  /** A LogicalType corresponding to TIME_WITH_LOCAL_TIME_ZONE. */

  public static class TimeWithLocalTzType extends LogicalTypes.PassThroughLogicalType<Instant> {

    public static final String IDENTIFIER = "SqlTimeWithLocalTzType";



    public TimeWithLocalTzType() {

      super(IDENTIFIER, "", Schema.FieldType.DATETIME);

    }

  }



  /** A LogicalType corresponding to TIMESTAMP_WITH_LOCAL_TIME_ZONE. */

  public static class TimestampWithLocalTzType extends LogicalTypes.PassThroughLogicalType<Instant> {

    public static final String IDENTIFIER = "SqlTimestampWithLocalTzType";



    public TimestampWithLocalTzType() {

      super(IDENTIFIER, "", Schema.FieldType.DATETIME);

    }

  }



  /** A LogicalType corresponding to CHAR. */

  public static class CharType extends LogicalTypes.PassThroughLogicalType<String> {

    public static final String IDENTIFIER = "SqlCharType";



    public CharType() {

      super(IDENTIFIER, "", Schema.FieldType.STRING);

    }

  }



  /** Returns true if the type is any of the various date time types. */

  public static boolean isDateTimeType(Schema.FieldType fieldType) {

    if (fieldType.getTypeName() == Schema.TypeName.DATETIME) {

      return true;

    }



    if (fieldType.getTypeName().isLogicalType()) {

      String logicalId = fieldType.getLogicalType().getIdentifier();

      return logicalId.equals(DateType.IDENTIFIER)

          || logicalId.equals(TimeType.IDENTIFIER)

          || logicalId.equals(TimeWithLocalTzType.IDENTIFIER)

          || logicalId.equals(TimestampWithLocalTzType.IDENTIFIER);

    }

    return false;

  }



  public static boolean isStringType(Schema.FieldType fieldType) {

    if (fieldType.getTypeName() == Schema.TypeName.STRING) {

      return true;

    }



    if (fieldType.getTypeName().isLogicalType()) {

      String logicalId = fieldType.getLogicalType().getIdentifier();

      return logicalId.equals(CharType.IDENTIFIER);

    }

    return false;

  }



  // The list of field type names used in SQL as Unified field types.

  public static final Schema.FieldType TINY_INT = Schema.FieldType.BYTE;

  public static final Schema.FieldType SMALL_INT = Schema.FieldType.INT16;

  public static final Schema.FieldType INTEGER = Schema.FieldType.INT32;

  public static final Schema.FieldType BIG_INT = Schema.FieldType.INT64;

  public static final Schema.FieldType FLOAT = Schema.FieldType.FLOAT;

  public static final Schema.FieldType DOUBLE = Schema.FieldType.DOUBLE;

  public static final Schema.FieldType DECIMAL = Schema.FieldType.DECIMAL;

  public static final Schema.FieldType BOOLEAN = Schema.FieldType.BOOLEAN;

  public static final Schema.FieldType VARBINARY = Schema.FieldType.BYTES;

  public static final Schema.FieldType VARCHAR = Schema.FieldType.STRING;

  public static final Schema.FieldType CHAR = Schema.FieldType.logicalType(new CharType());

  public static final Schema.FieldType DATE = Schema.FieldType.logicalType(new DateType());

  public static final Schema.FieldType TIME = Schema.FieldType.logicalType(new TimeType());

  public static final Schema.FieldType TIME_WITH_LOCAL_TZ =

      Schema.FieldType.logicalType(new TimeWithLocalTzType());

  public static final Schema.FieldType TIMESTAMP = Schema.FieldType.DATETIME;

  public static final Schema.FieldType TIMESTAMP_WITH_LOCAL_TZ =

      Schema.FieldType.logicalType(new TimestampWithLocalTzType());



  private static final BiMap<Schema.FieldType, SqlTypeName> BEAM_TO_CALCITE_TYPE_MAPPING =

      ImmutableBiMap.<Schema.FieldType, SqlTypeName>builder()

          .put(TINY_INT, SqlTypeName.TINYINT)

          .put(SMALL_INT, SqlTypeName.SMALLINT)

          .put(INTEGER, SqlTypeName.INTEGER)

          .put(BIG_INT, SqlTypeName.BIGINT)

          .put(FLOAT, SqlTypeName.FLOAT)

          .put(DOUBLE, SqlTypeName.DOUBLE)

          .put(DECIMAL, SqlTypeName.DECIMAL)

          .put(BOOLEAN, SqlTypeName.BOOLEAN)

          .put(VARBINARY, SqlTypeName.VARBINARY)

          .put(VARCHAR, SqlTypeName.VARCHAR)

          .put(CHAR, SqlTypeName.CHAR)

          .put(DATE, SqlTypeName.DATE)

          .put(TIME, SqlTypeName.TIME)

          .put(TIME_WITH_LOCAL_TZ, SqlTypeName.TIME_WITH_LOCAL_TIME_ZONE)

          .put(TIMESTAMP, SqlTypeName.TIMESTAMP)

          .put(TIMESTAMP_WITH_LOCAL_TZ, SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE)

          .build();



  private static final ImmutableMap<SqlTypeName, Schema.FieldType> CALCITE_TO_BEAM_TYPE_MAPPING =

      ImmutableMap.<SqlTypeName, Schema.FieldType>builder()

          .put(SqlTypeName.TINYINT, TINY_INT)

          .put(SqlTypeName.SMALLINT, SMALL_INT)

          .put(SqlTypeName.INTEGER, INTEGER)

          .put(SqlTypeName.BIGINT, BIG_INT)

          .put(SqlTypeName.FLOAT, FLOAT)

          .put(SqlTypeName.DOUBLE, DOUBLE)

          .put(SqlTypeName.DECIMAL, DECIMAL)

          .put(SqlTypeName.BOOLEAN, BOOLEAN)

          .put(SqlTypeName.VARBINARY, VARBINARY)

          .put(SqlTypeName.BINARY, VARBINARY)

          .put(SqlTypeName.VARCHAR, VARCHAR)

          .put(SqlTypeName.CHAR, CHAR)

          .put(SqlTypeName.DATE, DATE)

          .put(SqlTypeName.TIME, TIME)

          .put(SqlTypeName.TIME_WITH_LOCAL_TIME_ZONE, TIME_WITH_LOCAL_TZ)

          .put(SqlTypeName.TIMESTAMP, TIMESTAMP)

          .put(SqlTypeName.TIMESTAMP_WITH_LOCAL_TIME_ZONE, TIMESTAMP_WITH_LOCAL_TZ)

          .build();



  // Since there are multiple Calcite type that correspond to a single Unified type, this is the

  // default mapping.

  private static final Map<Schema.FieldType, SqlTypeName> BEAM_TO_CALCITE_DEFAULT_MAPPING =

      ImmutableMap.of(

          Schema.FieldType.DATETIME, SqlTypeName.TIMESTAMP,

          Schema.FieldType.STRING, SqlTypeName.VARCHAR);



  /** Generate {@link Schema} from {@code RelDataType} which is used to create table. */

  public static Schema toSchema(RelDataType tableInfo) {

    return tableInfo.getFieldList().stream().map(CalciteUtils::toField).collect(Schema.toSchema());

  }



  public static SqlTypeName toSqlTypeName(Schema.FieldType type) {

    switch (type.getTypeName()) {

      case ROW:

        return SqlTypeName.ROW;

      case ARRAY:

        return SqlTypeName.ARRAY;

      case MAP:

        return SqlTypeName.MAP;

      default:

        SqlTypeName typeName = BEAM_TO_CALCITE_TYPE_MAPPING.get(type.withNullable(false));

        if (typeName != null) {

          return typeName;

        } else {

          // This will happen e.g. if looking up a STRING type, and metadata isn't set to say which

          // type of SQL string we want. In this case, use the default mapping.

          return BEAM_TO_CALCITE_DEFAULT_MAPPING.get(type);

        }

    }

  }



  public static Schema.FieldType toFieldType(SqlTypeName sqlTypeName) {

    switch (sqlTypeName) {

      case MAP:

      case MULTISET:

      case ARRAY:

      case ROW:

        throw new IllegalArgumentException(

            String.format(

                "%s is a type constructor that takes parameters, not a type,"

                    + "so it cannot be converted to a %s",

                sqlTypeName, Schema.FieldType.class.getSimpleName()));

      default:

        return CALCITE_TO_BEAM_TYPE_MAPPING.get(sqlTypeName);

    }

  }



  public static Schema.Field toField(RelDataTypeField calciteField) {

    return toField(calciteField.getName(), calciteField.getType());

  }



  public static Schema.Field toField(String name, RelDataType calciteType) {

    return Schema.Field.of(name, toFieldType(calciteType)).withNullable(calciteType.isNullable());

  }



  public static Schema.FieldType toFieldType(RelDataType calciteType) {

    switch (calciteType.getSqlTypeName()) {

      case ARRAY:

      case MULTISET:

        return Schema.FieldType.array(toFieldType(calciteType.getComponentType()));

      case MAP:

        return Schema.FieldType.map(

            toFieldType(calciteType.getKeyType()), toFieldType(calciteType.getValueType()));

      case ROW:

        return Schema.FieldType.row(toSchema(calciteType));



      default:

        return toFieldType(calciteType.getSqlTypeName());

    }

  }



  /** Create an instance of {@code RelDataType} so it can be used to create a table. */

  public static RelDataType toCalciteRowType(Schema schema, RelDataTypeFactory dataTypeFactory) {

    RelDataTypeFactory.Builder builder = new RelDataTypeFactory.Builder(dataTypeFactory);



    IntStream.range(0, schema.getFieldCount())

        .forEach(

            idx ->

                builder.add(

                    schema.getField(idx).getName(), toRelDataType(dataTypeFactory, schema, idx)));

    return builder.build();

  }



  public static RelDataType toRelDataType(RelDataTypeFactory dataTypeFactory, Schema.FieldType fieldType) {

    switch (fieldType.getTypeName()) {

      case ARRAY:

        return dataTypeFactory.createArrayType(

            toRelDataType(dataTypeFactory, fieldType.getCollectionElementType()),

            UNLIMITED_ARRAY_SIZE);

      case MAP:

        RelDataType componentKeyType = toRelDataType(dataTypeFactory, fieldType.getMapKeyType());

        RelDataType componentValueType =

            toRelDataType(dataTypeFactory, fieldType.getMapValueType());

        return dataTypeFactory.createMapType(componentKeyType, componentValueType);

      case ROW:

        return toCalciteRowType(fieldType.getRowSchema(), dataTypeFactory);

      default:

        return dataTypeFactory.createSqlType(toSqlTypeName(fieldType));

    }

  }



  private static RelDataType toRelDataType(

	  RelDataTypeFactory dataTypeFactory, Schema schema, int fieldIndex) {

    Schema.Field field = schema.getField(fieldIndex);

    RelDataType type = toRelDataType(dataTypeFactory, field.getType());



    return dataTypeFactory.createTypeWithNullability(type, field.getType().getNullable());

  }



  /**

   * SQL-Java type mapping, with specified Unified rules: <br>

   * 1. redirect {@link AbstractInstant} to {@link Date} so Calcite can recognize it.

   *

   * @param rawType

   * @return

   */

  public static RelDataType sqlTypeWithAutoCast(RelDataTypeFactory typeFactory, Type rawType) {

    // For Joda time types, return SQL type for java.util.Date.

    if (rawType instanceof Class && AbstractInstant.class.isAssignableFrom((Class<?>) rawType)) {

      return typeFactory.createJavaType(Date.class);

    } else if (rawType instanceof Class && ByteString.class.isAssignableFrom((Class<?>) rawType)) {

      return typeFactory.createJavaType(byte[].class);

    }

    return typeFactory.createJavaType((Class) rawType);

  }

}